package com.pine.app.module.security.oauth.provider.token.store;

import com.pine.app.module.security.oauth.provider.OAuth2Authentication;
import com.pine.app.module.security.oauth.provider.token.OAuth2AccessToken;

import java.util.Date;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.DelayQueue;
import java.util.concurrent.Delayed;
import java.util.concurrent.TimeUnit;

/**
 * @author xiaoyuan
 * @create 2019/10/8 14:26
 **/
public class InMemoryTokenStore extends AbstractTokenStore {

    private final ConcurrentHashMap<String, OAuth2AccessToken> accessTokenStore = new ConcurrentHashMap();
    private final ConcurrentHashMap<String, OAuth2AccessToken> authenticationToAccessTokenStore = new ConcurrentHashMap();
    private final ConcurrentHashMap<String, OAuth2Authentication> authenticationStore = new ConcurrentHashMap();
    private final ConcurrentHashMap<String, OAuth2AccessToken> refreshTokenStore = new ConcurrentHashMap();
    private final ConcurrentHashMap<String, OAuth2Authentication> refreshTokenAuthenticationStore = new ConcurrentHashMap();
    private final DelayQueue<InMemoryTokenStore.TokenExpiry> expiryQueue = new DelayQueue();


    @Override
    public OAuth2Authentication readAuthentication(OAuth2AccessToken oAuth2AccessToken) {
        return this.authenticationStore.get(oAuth2AccessToken.getAccessToken());
    }

    @Override
    public OAuth2Authentication readAuthentication(String key) {
        return this.authenticationStore.get(key);
    }

    @Override
    public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication) {
//        if(this.flushCounter.incrementAndGet() >= this.flushInterval) {
//            this.flush();
//            this.flushCounter.set(0);
//        }
        this.accessTokenStore.put(token.getAccessToken(), token);
        this.authenticationStore.put(token.getAccessToken(), authentication);
        this.authenticationToAccessTokenStore.put(this.authenticationKeyGenerator.extractKey(authentication), token);
//        if(token.getExpireTime() != null) {
//            InMemoryTokenStore.TokenExpiry expiry = new InMemoryTokenStore.TokenExpiry(token.getAccessToken(), token.getExpireTime());
//            this.expiryQueue.remove(this.expiryMap.put(token.getAccessToken(), expiry));
//            this.expiryQueue.put(expiry);
//        }
        if(token.getOauth2RefreshToken()!=null){
             this.refreshTokenStore.put(token.getRefreshToken(),token);
            this.refreshTokenAuthenticationStore.put(token.getRefreshToken(),authentication);
        }
    }

    @Override
    public OAuth2AccessToken readAccessToken(String token) {
        return this.accessTokenStore.get(token);
    }

    @Override
    public void removeAccessToken(OAuth2AccessToken token) {
        this.removeAccessToken(token.getAccessToken());
    }

    @Override
    public void removeAccessToken(String token) {
        OAuth2AccessToken accessToken = this.accessTokenStore.get(token);
        if(accessToken!=null){
            this.accessTokenStore.remove(token);
            if(accessToken.getOauth2RefreshToken()!=null){
                this.refreshTokenStore.remove(accessToken.getRefreshToken());
                this.refreshTokenAuthenticationStore.remove(accessToken.getRefreshToken());
            }

        }
        OAuth2Authentication oAuth2Authentication = this.authenticationStore.get(token);
        if (oAuth2Authentication != null) {
            this.authenticationToAccessTokenStore.remove(this.authenticationKeyGenerator.extractKey(oAuth2Authentication));
        }
        this.authenticationStore.remove(token);
    }

    @Override
    public void removeRefreshToken(String refreshToken) {
        refreshTokenStore.remove(refreshToken);
        this.refreshTokenAuthenticationStore.remove(refreshToken);
    }

    @Override
    public OAuth2AccessToken readRefreshToken(String refreshToken) {
        return refreshTokenStore.get(refreshToken);
    }

    @Override
    public OAuth2Authentication readRefreshToAuthentication(String refreshToken) {
        return refreshTokenAuthenticationStore.get(refreshToken);
    }

    public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {
        String key = this.authenticationKeyGenerator.extractKey(authentication);
        OAuth2AccessToken accessToken = this.authenticationToAccessTokenStore.get(key);
        if (accessToken != null && !key.equals(this.authenticationKeyGenerator.extractKey(this.readAuthentication(accessToken.getAccessToken())))) {
            this.storeAccessToken(accessToken, authentication);
        }
        return accessToken;
    }

    private void flush() {
        for(InMemoryTokenStore.TokenExpiry expiry = this.expiryQueue.poll(); expiry != null; expiry = this.expiryQueue.poll()) {
            this.removeAccessToken(expiry.getValue());
        }

    }


    private static class TokenExpiry implements Delayed {
        private final long expiry;
        private final String value;

        public TokenExpiry(String value, Date date) {
            this.value = value;
            this.expiry = date.getTime();
        }

        public int compareTo(Delayed other) {
            if(this == other) {
                return 0;
            } else {
                long diff = this.getDelay(TimeUnit.MILLISECONDS) - other.getDelay(TimeUnit.MILLISECONDS);
                return diff == 0L?0:(diff < 0L?-1:1);
            }
        }

        public long getDelay(TimeUnit unit) {
            return this.expiry - System.currentTimeMillis();
        }

        public String getValue() {
            return this.value;
        }
    }
}
